from housegan_utils import *

if __name__=="__main__": # python housegan_train_probe.py --n_cpu 8
    parser = argparse.ArgumentParser()
    parser.add_argument("--n_epochs", type=int, default=1000000, help="number of epochs of training")
    parser.add_argument("--batch_size", type=int, default=32, help="size of the batches")
    parser.add_argument("--g_lr", type=float, default=0.0001, help="adam: learning rate")
    parser.add_argument("--d_lr", type=float, default=0.0001, help="adam: learning rate")
    parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
    parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
    parser.add_argument("--n_cpu", type=int, default=8, help="number of cpu threads to use during batch generation")
    parser.add_argument("--latent_dim", type=int, default=128, help="dimensionality of the latent space")
    parser.add_argument("--img_size", type=int, default=32, help="size of each image dimension")
    parser.add_argument("--sample_interval", type=int, default=50000, help="interval between image sampling")
    parser.add_argument("--exp_folder", type=str, default='exp', help="destination folder")
    parser.add_argument("--n_critic", type=int, default=1, help="number of training steps for discriminator per iter")
    parser.add_argument("--target_set", type=str, default='A', help="which split to remove")

    parser.add_argument("--gpu_id", type=str, default="0", help="gpu id")
    opt = parser.parse_args()
    print(opt)

    os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpu_id
    print('os.environ["CUDA_VISIBLE_DEVICES"]',os.environ["CUDA_VISIBLE_DEVICES"])
    device =torch.device("cuda:0")
    #device = torch.device("cuda:{}".format(os.environ["CUDA_VISIBLE_DEVICES"]))

    # checkpoint
    checkpoint_generator = "./housegan/checkpoints/generator_exp_example_D_500000.pth"
    checkpoint_discriminator = "./housegan/checkpoints/discriminator_exp_example_D_500000.pth"
    # init models
    print("models init and load checkpoints...")
    generator = Generator()
    discriminator = Discriminator()

    generator.to(device)
    discriminator.to(device)
    generator.load_state_dict(torch.load(checkpoint_generator))
    discriminator.load_state_dict(torch.load(checkpoint_discriminator))
    print("load datasets...")
    rooms_path = './housegan/dataset'#'/home/nelson/Workspace/autodesk/housegan/'
    fp_dataset_train = FloorplanGraphDataset(rooms_path, transforms.Normalize(mean=[0.5], std=[0.5]), target_set=opt.target_set)
    #fp_subdataset_train = torch.utils.data.Subset(fp_dataset_train, range(0, len(fp_dataset_train)//100))
    fp_dataloader_train = torch.utils.data.DataLoader(fp_dataset_train, 
                                            batch_size=64,#opt.batch_size, 
                                            shuffle=True,
                                            num_workers=opt.n_cpu,
                                            collate_fn=floorplan_collate_fn,
                                            pin_memory = True)
    fp_dataset_test = FloorplanGraphDataset(rooms_path, transforms.Normalize(mean=[0.5], std=[0.5]), target_set=opt.target_set, split='eval')
    fp_dataloader_val = torch.utils.data.DataLoader(fp_dataset_test, 
                                            batch_size=64,#opt.batch_size, 
                                            shuffle=True,
                                            num_workers=opt.n_cpu,
                                            collate_fn=floorplan_collate_fn,
                                            pin_memory = True)
    
    # 使用未训练的生成器来生成负样本
    #generator = Generator().to(device)
    generator.eval()
    discriminator.eval()

    # 为每层的 features 分别初始化一个 LinearClassifierProbe, 并设置 optimizer
    print("init probes...")
    probes = dict()
    probe_optimizers = dict()
    learning_rate =3e-4
    n_epochs=10000
    decrease_domain_length_in_valid=20

    for feature_name, feature_length in DISCRIMINATOR_MIDFEATURE_SIZE.items():
        probes[feature_name] = LinearClassifierProbe(n_in_features=feature_length, n_out_features=2, softmax_dim=-1)
        probes[feature_name].to(device)
        probes[feature_name].train()
        probe_optimizers[feature_name] = torch.optim.Adam(probes[feature_name].parameters(), lr=learning_rate)
    print("training probes...")
    print(f"lr = {learning_rate}")
    print(f"n_epochs = {n_epochs}")
    print(f"decrease_domain_length_in_valid = {decrease_domain_length_in_valid}")
    
    loss_epochs, error_epochs, probe_stop = train_probes(opt, generator, discriminator,fp_dataloader_train, fp_dataloader_val,\
        probes, probe_optimizers, n_epochs=n_epochs, decrease_domain_length_in_valid = decrease_domain_length_in_valid)
    os.makedirs("housegan_probe_training_result",exist_ok=True)
    with open("./housegan_probe_training_result/probe_loss_and_error_last.pkl","wb") as f:
        pickle.dump([loss_epochs, error_epochs, probe_stop],f)